#include <random>
#include <cstring>
#include <chrono>
#include <algorithm>
#include <fstream>
#include <sstream>
#include <thread>

#include <caffe/net.hpp>
#include <caffe/profiler.hpp>
#include <opencv2/opencv.hpp>

using namespace std;
using namespace caffe;

#include <windows.h> 
/*! \brief Timer */
class Timer {
	using Clock = std::chrono::high_resolution_clock;
public:
	/*! \brief start or restart timer */
	inline void Tic() {
		start_ = Clock::now();
	}
	/*! \brief stop timer */
	inline void Toc() {
		end_ = Clock::now();
	}
	/*! \brief return time in ms */
	inline double Elasped() {
		auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_ - start_);
		return duration.count();
	}
private:
	Clock::time_point start_, end_;
};


void WrapInputLayer(const cv::Mat& img, std::vector<cv::Mat> *input_channels, std::shared_ptr<caffe::Net> cnn) {
	shared_ptr<Blob> input_layer = cnn->blob_by_name("data");
	int width = img.cols;
	int height = img.rows;
	int channel = img.channels();
	int batchsize = 1;
	input_layer->Reshape(batchsize, channel, height, width);
	float* input_data = input_layer->mutable_cpu_data();
	for (int j = 0; j < input_layer->channels(); ++j) {
		cv::Mat chennel(height, width, CV_32FC1, input_data);
		input_channels->push_back(chennel);
		input_data += width * height;
	}
	cv::split(img, *input_channels);
}

void print_shape(const std::vector<int>& shape) {
	for (int i = 0; i < shape.size(); i++) {
		printf("%d ", shape[i]);
	}
	printf("\n");
}
void Forward(const cv::Mat &data, std::shared_ptr<caffe::Net> cnn) {
	assert(data.channels() == 3);
	std::vector<cv::Mat> input_channels;
	WrapInputLayer(data, &input_channels, cnn);
	cnn->Forward();

	shared_ptr<Blob> net_output = cnn->blob_by_name("net_output");
	print_shape(net_output->shape());
}

int main() {

	TCHAR NPath[MAX_PATH];
	GetCurrentDirectory(MAX_PATH, NPath);
	std::cout << "current directory is " << NPath << "\n";
	if (caffe::GPUAvailable()) {
		printf("GPU is available\n");
		caffe::SetMode(caffe::GPU, 0);
	}
	else {
		printf("Only CPU is available\n");
	}

	std::string prototxt = "../models/pose_deploy_linevec.prototxt";
	std::string caffemodel = "../models/pose_iter_440000.caffemodel";

	std::shared_ptr<caffe::Net> ssh = std::make_shared<caffe::Net>(prototxt);
	ssh->CopyTrainedLayersFrom(caffemodel);

	cv::Mat image = cv::imread("../misc/test_imgs/00.jpg");
	cv::resize(image, image, cv::Size(224, 224));

	//ssh->Reshape()

	Forward(image, ssh);

	return 0;
}